package org.skipper.ai.strategy.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import lombok.extern.slf4j.Slf4j;
import org.skipper.ai.strategy.ChatStrategy;
import org.skipper.ai.strategy.domain.ChatParam;
import org.skipper.ai.strategy.domain.aliyun.AliyunAiProperties;
import org.skipper.ai.strategy.domain.aliyun.AliyunChatParam;
import org.skipper.ai.strategy.result.ChatResult;
import org.skipper.ai.strategy.result.aliyun.AliyunChatResult;
import org.skipper.ai.utils.SpringContextUtil;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;

/**
 * Created with IntelliJ IDEA.
 * Description: TODO
 *
 * @author shoujian
 * Date: 2025-10-09
 * Time: 16:43
 */
@Slf4j
public class AliyunDashScopeChat implements ChatStrategy {

    // 若使用新加坡地域的模型，请释放下列注释
    // static {
    //     Constants.baseHttpApiUrl="https://dashscope-intl.aliyuncs.com/api/v1";
    // }
    public static GenerationParam createGenerationParam(List<Message> messages, String model,
                                                        Boolean enableThinking) {
        AliyunAiProperties properties = SpringContextUtil.getBean(AliyunAiProperties.class);

        return GenerationParam.builder()
                // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                // 新加坡和北京地域的API Key不同。获取API Key：https://help.aliyun.com/zh/model-studio/get-api-key
                .apiKey(properties.getApiKey())
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model(model)
                .enableThinking(enableThinking)
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
    }

    public static GenerationResult callGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.call(param);
    }

    private static Message createMessage(String role, String content) {
        return Message.builder().role(role).content(content).build();
    }

    @Override
    public ChatResult chat(ChatParam chatParam) {

        if (chatParam instanceof AliyunChatParam) {
            AliyunChatParam aliyunChatParam = (AliyunChatParam) chatParam;
            try {
                List<Message> messages = aliyunChatParam.getAliyunMessages().stream()
                        .map(message -> createMessage(message.getRole(), message.getContent()))
                        .collect(Collectors.toList());
                GenerationParam param = createGenerationParam(messages, aliyunChatParam.getModel(), aliyunChatParam.getEnableThinking());
                GenerationResult result = callGenerationWithMessages(param);

                log.info("GenerationResult: " + result);

                AliyunChatResult aliyunChatResult = new AliyunChatResult();
                aliyunChatResult.setInputTokens(result.getUsage().getInputTokens());
                aliyunChatResult.setOutputTokens(result.getUsage().getOutputTokens());
                aliyunChatResult.setText(result.getOutput().getChoices().get(0).getMessage().getContent());
                aliyunChatResult.setTotalTokens(result.getUsage().getTotalTokens());
                return aliyunChatResult;
            } catch (ApiException | NoApiKeyException | InputRequiredException e) {
                e.printStackTrace();
            }

        } else {
            throw new IllegalArgumentException("参数类型错误");
        }

        return null;
    }
}
